import os
import argparse
import numpy as np

def _replace(path):
    path = path.replace("audio", "wav")
    path = path.replace("flac", "wav")
    return path

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_dir', help='dataset dir', type=str, default="SITW")
    parser.add_argument('--output_trial_dir', help='output trial path', type=str, default="data/")
    args = parser.parse_args()

    for sub in ["eval", "dev"]:
        enroll_lst_path = os.path.join(args.dataset_dir, "{}/lists/enroll-core.lst".format(sub))
        raw_trial_path = os.path.join(args.dataset_dir, "{}/keys/core-core.lst".format(sub))

        spk2wav_mapping = {}
        enroll_lst = np.loadtxt(enroll_lst_path, str)
        for item in enroll_lst:
            spk2wav_mapping[item[0]] = item[1]
        trials = np.loadtxt(raw_trial_path, str)

        output_trial_path = os.path.join(args.output_trial_dir, "SITW_{}.lst".format(sub))
        with open(output_trial_path, "w") as f:
            for item in trials:
                enroll_path = os.path.join(args.dataset_dir, sub, spk2wav_mapping[item[0]])
                test_path = os.path.join(args.dataset_dir, sub, item[1])
                label = item[2]
                if label == "tgt":
                    label = "1"
                else:
                    label = "0"

                enroll_path = _replace(enroll_path)
                test_path = _replace(test_path)

                f.write("{} {} {}\n".format(label, enroll_path, test_path))

